# -*- coding: utf-8 -*-
"""
ERA5数据加载工具模块

此模块提供用于加载和处理ERA5气象数据的类和函数。
主要功能包括：
- 加载ERA5 NetCDF格式数据
- 提取指定时间段和区域的数据
- 数据预处理和格式转换
- 支持多种数据格式和时间维度

作者: Jia
日期: 2025-06-20
"""

import os
import gc

import numpy as np
from typing import Dict, List, Optional, Tuple, Any
from dataclasses import dataclass

from ..config import get_config, ROOT_DIR
from ..parsers import create_parser


@dataclass
class Point:
    """表示地理坐标点的数据类

    提供统一的接口来表示和处理地理坐标点，支持多种格式的转换。
    """
    longitude: float  # 经度
    latitude: float  # 纬度

    def __str__(self) -> str:
        """返回点的字符串表示，格式为 'longitude_latitude'"""
        return f"{self.longitude:.4f}_{self.latitude:.4f}"

    def to_tuple(self) -> Tuple[float, float]:
        """返回点的元组表示 (longitude, latitude)"""
        return self.longitude, self.latitude

    def to_dict(self) -> Dict[str, float]:
        """返回点的字典表示 {'longitude': longitude, 'latitude': latitude}"""
        return {"longitude": self.longitude, "latitude": self.latitude}

    @classmethod
    def from_tuple(cls, coords: Tuple[float, float]) -> 'Point':
        """从元组创建点对象"""
        return cls(longitude=coords[0], latitude=coords[1])

    @classmethod
    def from_array(cls, coords: np.ndarray) -> 'Point':
        """从numpy数组创建点对象"""
        return cls(longitude=float(coords[0]), latitude=float(coords[1]))

    @classmethod
    def from_object(cls, obj: Any) -> Optional['Point']:
        """从各种对象创建点对象"""
        # 避免循环引用，直接在这里实现坐标提取逻辑
        if isinstance(obj, Point):
            return obj
        elif isinstance(obj, tuple) and len(obj) == 2:
            return cls(longitude=obj[0], latitude=obj[1])
        elif isinstance(obj, np.ndarray) and len(obj) == 2:
            return cls(longitude=float(obj[0]), latitude=float(obj[1]))
        elif hasattr(obj, 'longitude') and hasattr(obj, 'latitude'):
            return cls(longitude=obj.longitude, latitude=obj.latitude)
        elif hasattr(obj, 'lon') and hasattr(obj, 'lat'):
            return cls(longitude=obj.lon, latitude=obj.lat)
        elif hasattr(obj, '经度') and hasattr(obj, '纬度'):
            return cls(longitude=obj['经度'], latitude=obj['纬度'])
        else:
            print(f"警告: 无法从{obj}中提取经纬度信息，跳过该点")
            return None


class ERA5DataLoader:
    """
    ERA5数据加载器类
    
    提供统一的接口来加载和处理ERA5气象数据，支持多种数据格式和时间维度。
    """

    def __init__(self, data_dir: Optional[str] = None):
        """
        初始化数据加载器
        
        Args:
            data_dir: ERA5数据目录，如为None则使用默认目录
        """
        # 设置数据目录
        self.data_dir = data_dir
        if data_dir is None:
            self.data_dir = get_config('era5_data_path', 'Daily_ERA5(1960-2024)')

        # 检查数据目录是否为绝对路径
        if not os.path.isabs(self.data_dir):
            # 尝试获取项目根目录
            try:
                # 构建相对于项目根目录的绝对路径
                abs_path = os.path.join(ROOT_DIR, self.data_dir)
                if os.path.exists(abs_path):
                    self.data_dir = abs_path
                    print(f"使用项目根目录相对路径: {self.data_dir}")
            except ImportError:
                print("无法导入ROOT_DIR，使用相对路径")

        # 打印初始化信息
        print(f"初始化ERA5数据加载器:")
        print(f"  - 数据目录: {self.data_dir}")
        print(f"  - 数据目录是否存在: {os.path.exists(self.data_dir)}")

    def load_era5_data(self, start_year: int, end_year: int, months: List[int],
                       points: List[np.ndarray], chunk_size: Optional[int] = None) -> Dict:
        """
        加载ERA5降水数据，优化内存使用
        
        Args:
            start_year: 起始年份
            end_year: 结束年份
            months: 月份列表
            points: 经纬度点位列表（必选），格式为[(lon1, lat1), (lon2, lat2), ...]、numpy数组或其他包含经纬度信息的对象列表
            chunk_size: 一次处理的年份数量，默认为None（不分块处理）
            
        Returns:
            数据字典，格式为 {point_id: {year: {"total_precip": float}}}
            其中point_id为经纬度字符串（如"120.000000_30.000000"），year为年份，total_precip为该点位该年份的总降水量
        """
        # 检查数据目录是否存在
        assert os.path.exists(
            self.data_dir), f"警告: ERA5数据目录不存在: {self.data_dir}，请在配置文件中设置正确的era5_data_path路径"

        # 检查points参数
        assert points is not None and len(points) > 0, "警告: points参数不能为空，必须提供至少一个点位"
        print(f"批量加载{len(points)}个点位的数据")

        # 如果不指定chunk_size，直接加载所有数据
        if chunk_size is None:
            return self._load_points_data(start_year, end_year, months, points)

        # 按chunk_size分组处理年份
        result = {}
        for chunk_start in range(start_year, end_year + 1, chunk_size):
            chunk_end = min(chunk_start + chunk_size - 1, end_year)
            print(f"处理年份区间: {chunk_start}-{chunk_end}")

            # 加载当前chunk的数据
            chunk_data = self._load_points_data(chunk_start, chunk_end, months, points)

            # 合并年份数据，而不是简单地更新字典
            for point_id, year_data in chunk_data.items():
                if point_id not in result:
                    result[point_id] = {}
                # 合并年份数据
                result[point_id].update(year_data)
            print(f"已加载{chunk_end - chunk_start + 1}个年份数据")
            # 每处理完一个chunk，主动触发垃圾回收
            gc.collect()

        return result

    def _load_points_data(self, start_year: int, end_year: int, months: List[int], points: List[Any],
                          ) -> dict[str, dict[Any, Any]]:
        """
        从ERA5数据中提取指定年份、月份和点位的降水数据

        Args:
            start_year: 起始年份
            end_year: 结束年份
            months: 月份列表
            points: 经纬度点位列表（必选），可以是Point对象、元组、numpy数组或其他包含经纬度信息的对象列表

        Returns:
            数据字典，格式为 {point_id: {year: {"total_precip": float}}}
            其中point_id为经纬度字符串（如"120.000000_30.000000"），year为年份，total_precip为该点位该年份的总降水量
        """
        assert points is not None and len(points) > 0, "警告: points参数不能为空，必须提供至少一个点位"
        print(f"开始收集{start_year}-{end_year}年所有区域{months}数据...")

        # 初始化数据结构
        grid_data = {}
        for year in range(start_year, end_year + 1):
            # 构建文件路径
            file_path = f"{self.data_dir}/ERA5.P.daily.{year}.nc"
            # 直接使用xarray打开数据集
            parser = create_parser(file_path) if create_parser else None

            # 获取时间维度名称
            time_dim = parser.get_time_dimension()
            if time_dim is None:
                raise IndexError(f"警告: 未找到标准时间维度，使用第一个维度作为时间维度")

            # 提取变量数据
            var_data = parser.extract_variable()

            # 获取时间值
            dates = [parser.format_date(t) for t in var_data[time_dim].values]
            data_months = [int(date.split('-')[1]) for date in dates]  # 提取月份
            # 提取指定月份的日期 index
            month_indices = []
            for month in months:
                month_indices.extend([i for i, m in enumerate(data_months) if m == month])
            if len(month_indices) < 1:
                print(f"警告: {year}年{months}月的数据为空")
                return grid_data
            # 提取数据
            month_data = var_data.isel({time_dim: month_indices})
            # 获取经纬度维度名称
            lat_dim, lon_dim = parser.get_spatial_dimensions()
            # 为每个点创建数据结构
            for i, point_obj in enumerate(points):
                # 尝试转换为Point对象
                point = Point.from_object(point_obj) if not isinstance(point_obj, Point) else point_obj
                if point is None:
                    continue
                point_lon, point_lat = point.longitude, point.latitude
                # 生成点位ID
                point_id = str(point)  # 使用Point的__str__方法生成ID
                # 初始化网格点数据结构
                if point_id not in grid_data:
                    grid_data[point_id] = {}

                # 找到最接近的网格点索引
                lat_idx = abs(month_data[lat_dim] - point_lat).argmin().item()
                lon_idx = abs(month_data[lon_dim] - point_lon).argmin().item()
                # 使用isel方法直接提取点位数据
                point_data_array = month_data.isel({lat_dim: lat_idx, lon_dim: lon_idx}).values
                total_precip = np.sum(point_data_array)
                # 存储数据
                grid_data[point_id][year] = {
                    # 'precip_values':  point_data_array.tolist(),
                    'total_precip': float(total_precip)
                }
            print(f"已加载{year}年的{len(points)}个点位数据")
        return grid_data


def parse_geo_coordinates(point) -> tuple | None:
    """
    从各种格式的点位对象中解析提取经纬度坐标

    Args:
        point: 点位对象，可以是元组、numpy数组、Point类实例或其他包含经纬度信息的对象

    Returns:
        经纬度坐标元组 (longitude, latitude)，如果无法提取则返回 None
    """
    if isinstance(point, Point):
        return point.longitude, point.latitude
    elif isinstance(point, tuple) and len(point) == 2:
        return point
    elif isinstance(point, np.ndarray) and len(point) == 2:
        return point[0], point[1]
    elif hasattr(point, 'longitude') and hasattr(point, 'latitude'):
        return point.longitude, point.latitude
    elif hasattr(point, 'lon') and hasattr(point, 'lat'):
        return point.lon, point.lat
    elif hasattr(point, '经度') and hasattr(point, '纬度'):
        return point['经度'], point['纬度']
    else:
        print(f"警告: 无法从{point}中提取经纬度信息，跳过该点")
        return None


# 向后兼容的函数接口
def load_era5_data(data_dir: str, start_year: int, end_year: int,
                   months: List[int], points: List[np.ndarray], chunk_size: Optional[int] = None) -> Dict:
    """向后兼容的ERA5数据加载函数
    
    Args:
        data_dir: ERA5数据目录路径
        start_year: 开始年份
        end_year: 结束年份
        months: 要加载的月份列表
        points: 经纬度点位列表（必选），格式为[(lon1, lat1), (lon2, lat2), ...]、numpy数组或其他包含经纬度信息的对象列表
        chunk_size: 每次处理的年份数量，默认为None（不分块处理）
                   如果指定，则按指定大小分块处理年份数据
    
    Returns:
        ERA5数据字典，格式为 {point_id: {year: {"total_precip": float}}}
        其中 point_id为经纬度字符串（如"120.000000_30.000000"），year为年份，total_precip为该点位该年份的总降水量
    """
    assert os.path.exists(data_dir), f"ERA5数据目录不存在: {data_dir}"
    assert points is not None and len(points) > 0, "警告: points参数不能为空，必须提供至少一个点位"

    # 打印数据加载信息
    print(f"加载ERA5数据:")
    print(f"  - 数据目录: {data_dir}")
    print(f"  - 年份范围: {start_year}-{end_year}")
    print(f"  - 月份: {months}")
    print(f"  - 批量加载点位数量: {len(points)}")
    if chunk_size is None:
        print(f"  - 不分块处理")
    else:
        print(f"  - 分块大小: {chunk_size}")

    loader = ERA5DataLoader(data_dir)
    result = loader.load_era5_data(start_year, end_year, months, points, chunk_size=chunk_size)

    # 检查结果
    if len(result) == 0:
        print("警告: 加载的ERA5数据为空")
    else:
        print(f"成功加载ERA5数据，包含{len(result)}个点位")

    return result
